from Trader import Trader
from machine_learning import train_model
import math

#机器学习策略类，继承自交易者类
class MLTrader(Trader):
    def __init__(self, initial_money=0,  tax_ratio=0.001, max_stock_num = 10,stocks_info = {}):
        Trader.__init__(self, initial_money,  tax_ratio, max_stock_num , stocks_info)
        #用历史数据训练一个model
        self.model = train_model(stocks_info)
    def strategy(self, date):
        #选取几个显著的因子作为feature
        index = [23, 21, 24, 33, 6, 32, 26, 19, 25, 18, 27, 20]
        up_stocks = []
        down_stocks = []
        #卖出所有预测为下跌的股票
        for stock in self.stocks_info:
            stock_feature = [[self.stocks_info[stock][date].values[i] for i in index]]
            if self.model.predict(stock_feature)[0]==0 and not math.isnan(self.stocks_info[stock][date]['close']):
                down_stocks.append(stock)
        for stock in down_stocks:
            self.sell(date, stock_name=stock, share=self.stock_shares[stock])
        #在允许购买的范围内，买入所有预测为上涨的股票
        buy_num = self.max_stock_num
        for stock in self.stock_shares:
            if self.stock_shares[stock]>0:
                buy_num-=1
        for stock in self.stocks_info:
            stock_feature = [[self.stocks_info[stock][date].values[i] for i in index]]
            if not math.isnan(self.stocks_info[stock][date]['close']) and self.model.predict(stock_feature)[0]==1:
                if self.stock_shares[stock]>0:
                    up_stocks.append(stock)
                elif buy_num>0:
                    buy_num-=1
                    up_stocks.append(stock)
        for stock in up_stocks:
            money = self.cash/len(up_stocks)
            self.buy(date, stock_name=stock, money=money)
            
        self.net_worth = self.get_worth(date)
